package mcts

import (
	"fmt"
	"math"
)

func UCB1(node *MCTSNode) float64 {
	// v / visits   + 探索值
	if node.Visits == 0 {
		return math.Inf(1)
	}
	if node.Parent == nil {
		return 0
	}

	explodedValue := 2 * math.Log(float64(node.Parent.Visits)/float64(node.Visits))
	return node.Value/float64(node.Visits) + math.Sqrt(explodedValue)
}

func MCTS(root *MCTSNode, iterations int) *MCTSNode {
	leaf := root
	var err error
	for i := 0; i < iterations; i++ {
		println(i, "start")
		leaf = Select(leaf)
		var child *MCTSNode
		if !leaf.Env.IsTerminal() {
			if len(leaf.UntriedActions) > 0 {
				child = Expand(leaf)
			} else {
				child, err = BestChild(leaf)
				if err != nil {
					println("this leaf has no untried actions or vaild children. dont try again")
					leaf.done = true
					leaf = leaf.Parent
				} else {
					child = child.Move(child.action)
				}
			}
		} else {
			child = leaf
		}
		result := Simulate(child)
		Backpropagate(child, result)

		// show
		ShowMCTSNode(leaf.Env)

		if root.Env.IsTerminal() {
			leaf = root
			root.Env.Reset()
		}

	}
	best, _ := BestChild(root)
	return best
}

func BestChild(node *MCTSNode) (*MCTSNode, error) {
	// add pos
	bestValue := math.Inf(-1)
	bestNode := &MCTSNode{}
	hasVaildChildren := false
	for i, child := range node.Children {
		if child.done {
			continue
		}
		ucb1Value := UCB1(child)
		if ucb1Value > bestValue {
			bestValue = ucb1Value
			bestNode = child
			bestNode.pos = fmt.Sprintf("%d,%d done:%v", child.depth, i, bestNode.done)
			hasVaildChildren = true
		}
	}

	if !hasVaildChildren {
		return node, fmt.Errorf("has not a vaild child")
	}
	return bestNode, nil
}

func Select(node *MCTSNode) *MCTSNode {
	currentNode := node
	var err error
	if !currentNode.IsLeaf() {
		currentNode, err = BestChild(currentNode)
		println(".. choose", currentNode.pos)
		if err != nil {
			println(err.Error())
		}
	}
	if currentNode.done {
		currentNode = currentNode.Parent
		println(".. choose due to done", currentNode.pos)
	}
	println("Select", currentNode.pos)
	currentNode.ShowTrace()
	return currentNode
}

func Expand(node *MCTSNode) *MCTSNode {
	println("Expand start")
	action := node.BestAction(node.UntriedActions)
	node.UntriedActions = node.UntriedActions[1:]
	child := node.Move(action)
	node.AddChild(child)

	println("Expand", child.pos, child.done)
	return child
}

func Simulate(node *MCTSNode) float64 {
	env := *node.Env
	println("Simulate", node.pos)
	for !env.IsTerminal() {
		actions := env.GetLegalActions()
		action := node.BestAction(actions)
		env.Move(action)
	}
	return env.GetResult()
}

func Backpropagate(node *MCTSNode, result float64) {
	println("Backpropagate", node.pos)
	currentNode := node
	for currentNode != nil {
		currentNode.Visits++
		currentNode.Value += result * float64(node.Player)
		currentNode = currentNode.Parent
	}
}
